Feat: add bounded whole-run FIFO admission - #1541
Conversation
📝 WalkthroughWalkthroughThe change introduces generation-safe pipeline leases, run-scoped FIFO admission, task-acceptance signaling, lease-aware mailbox dispatch, and AICore image-aware stream reuse across C++, Python bindings, device runners, documentation, and tests. ChangesPipeline admission and acceptance
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 9
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/common/hierarchical/orchestrator.cpp (1)
177-206: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winAvoid racing
fail_run_submissionagainst scheduler activation. The run phase is sampled without the completion/run lock, andcancel_unstarted_runis later allowed to mark slotsPENDING/READY/FREEasFAILED; if the closed run has already been activated toEXECUTING, the scheduler can dispatch the same slot asRUNNINGconcurrently, losing one terminal transition. Guard cancellation by the current run phase and use stateful atomic transitions for slot states.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/common/hierarchical/orchestrator.cpp` around lines 177 - 206, Update fail_run_submission and cancel_unstarted_run to synchronize phase inspection with the run/completion state, preventing cancellation from racing scheduler activation. Recheck the current phase while holding the appropriate lock before cancelling, and make each slot-state update use compare-and-exchange transitions that only convert valid non-running states; never mark an EXECUTING/RUNNING slot FAILED or overwrite a concurrent scheduler transition.
🧹 Nitpick comments (5)
src/common/platform/onboard/host/device_runner_base.h (1)
297-320: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDoc block missing the new
aicore_image_hashparam. The@paramlist enumerates every other argument; add a line for the new one so the staging contract stays self-describing (same forrecord_host_orch_callableat Line 332).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/common/platform/onboard/host/device_runner_base.h` around lines 297 - 320, Update the documentation for record_device_orch_callable to add an `@param` entry describing aicore_image_hash, matching its position and purpose in the function signature. Also add the corresponding parameter documentation to record_host_orch_callable, without changing implementation behavior.src/common/hierarchical/scheduler.cpp (1)
366-369: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueOptional: fold the
active_run_cb ? scoped : unscopedbranching into helpers. The same "resolve active run, bail on invalid, pick scoped vs unscoped queue op" shape is now repeated in all four dispatch paths plus the wait predicate; a pair of small private helpers (e.g.resolve_active_run()returningstd::optional<RunId>andpop_ready(...)) would keep future queue-API changes in one place.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/common/hierarchical/scheduler.cpp` around lines 366 - 369, Consolidate the repeated active-run resolution and scoped/unscoped ready-queue branching across all four dispatch paths and the wait predicate in the scheduler class. Add small private helpers such as resolve_active_run() and pop_ready(...) that preserve invalid-run early returns while centralizing queue API selection, then update the affected paths to use them.python/simpler/task_interface.py (1)
1277-1284: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd one-line docstrings to match the neighbouring properties.
Every other property in this class documents its unit/semantics;
pipeline_depthandruntime_slot_countare the odd ones out.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/simpler/task_interface.py` around lines 1277 - 1284, Add concise one-line docstrings to the pipeline_depth and runtime_slot_count properties in the task interface, matching the neighboring properties’ style and documenting each value’s unit or semantics.tests/ut/cpp/hierarchical/test_orchestrator.cpp (1)
656-665: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winA failed
ASSERT_*between here andthird.get()turns the test into a hang.
ASSERT_TRUEreturns from the test body, and~futurefromstd::asyncblocks until the task completes — butbegin_run()only returns once a slot frees, which is exactly what the skipped lines would have done. Consider releasing the admission slot from a scope guard so an assertion failure fails fast instead of wedging the suite.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/ut/cpp/hierarchical/test_orchestrator.cpp` around lines 656 - 665, Ensure the asynchronous third begin_run test always releases an admission slot before any ASSERT_* can exit the test, by adding a scope guard around the first run’s slot ownership and consumption cleanup. Preserve the existing successful-path behavior while making cleanup execute during assertion-induced early returns, so third.get() cannot leave the test hanging.src/common/platform/onboard/host/device_runner_base.cpp (1)
143-150: 🗄️ Data Integrity & Integration | 🔵 Trivial | ⚡ Quick winBank validation accepts
< PTO_PIPELINE_MAX_DEPTHbut only two banks exist.
select_arena_bankadmits any id belowPTO_PIPELINE_MAX_DEPTH, while every consumer is a binarybank == 0 ? bank0 : bank1. With today's depth of 2 this is benign, but raising the depth constant would silently alias banks 2..N onto bank1 and share arenas across concurrently admitted runs. Prefer indexing an array of banks, or validate against the actual bank count.♻️ Sketch
- if (bank_id >= PTO_PIPELINE_MAX_DEPTH) { - LOG_ERROR("arena bank %u is outside [0, %u)", bank_id, PTO_PIPELINE_MAX_DEPTH); + if (bank_id >= kArenaBankCount) { + LOG_ERROR("arena bank %u is outside [0, %u)", bank_id, kArenaBankCount); return -1; }with
gm_heap_arena_[kArenaBankCount]etc. so the selection isgm_heap_arena_[g_arena_bank].Also applies to: 193-211
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/common/platform/onboard/host/device_runner_base.cpp` around lines 143 - 150, Update select_arena_bank and the associated arena consumers to validate against the actual number of supported banks rather than PTO_PIPELINE_MAX_DEPTH. Define or reuse a two-bank count, reject IDs outside that range, and replace binary bank-selection logic with indexed access so each accepted bank maps to its own arena without aliasing.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/orchestrator.md`:
- Around line 101-106: The documentation around the “later submit” behavior must
distinguish admission from acceptance: clarify that acceptance unblocks only an
already-admitted successor, while a third submission with bounded depth two
remains blocked in begin_run() until a terminal run releases its lease. Preserve
the existing endpoint acceptance and completion fallback details.
In `@docs/worker-manager.md`:
- Around line 175-181: Update the documentation around the mailbox polling loop
to state that on_accept() is a one-shot callback, invoked only on the first
observation of TASK_ACCEPTED. Document the acceptance_observed guard and
preserve TASK_DONE as the mailbox ownership boundary, so implementations do not
repeatedly decrement per-run acceptance accounting.
In `@python/simpler/worker.py`:
- Line 1544: In the comment near the worker implementation, replace the
ambiguous multiplication symbol “×” in “N×40B” with the ASCII character “x”,
preserving the rest of the comment unchanged.
- Around line 180-186: Update _PIPELINE_LEASE_FMT to use the C ABI’s packed
12-byte layout for PipelineSlotLease: uint32 slot_id, uint32 reserved, and
uint64 generation with no padding between fields. Ensure _OFF_ARGS resolves to
the established MAILBOX_OFF_ARGS offset and remains compatible with the C++
definition in pto_runtime_c_api.h.
In `@src/a2a3/platform/onboard/host/device_runner.cpp`:
- Around line 621-624: Make acceptance durable across the mailbox transition in
the dispatch flow around publish_task_accepted and
LocalMailboxEndpoint::run_with_accept: ensure a parent that observes TASK_DONE
without previously sampling TASK_ACCEPTED still invokes on_accept exactly once.
Use a latched acceptance flag/counter or equivalent state, and preserve the
existing behavior for parents that observe TASK_ACCEPTED first so
decrement_run_accepts is not duplicated.
In `@src/common/hierarchical/orchestrator.cpp`:
- Around line 360-371: Update Orchestrator::decrement_run_accepts to acquire
run->completion_mu before changing pending_accepts and before calling
completion_cv.notify_all(). Keep the decrement, underflow handling, error
recording, and notification behavior intact while ensuring the predicate
transition is synchronized with wait_run_accepted.
In `@src/common/hierarchical/types.cpp`:
- Around line 70-82: Update ReadyQueue’s unscoped access paths, including
try_pop(TaskSlot&) and try_front(TaskSlot&), so they no longer select buckets by
unordered_map iteration order; remove these overloads or add separate
run-insertion-order tracking that preserves FIFO, and update callers such as
try_pop_group(TaskSlot&) and try_pop_single(worker_id, TaskSlot&) accordingly.
In `@src/common/hierarchical/worker_manager.cpp`:
- Around line 397-404: Move the acceptance callback out of the mailbox-locked
section in the worker task flow around read_mailbox_state and
LocalMailboxEndpoint::run. Record that TASK_ACCEPTED was observed while polling,
then invoke on_accept_ only after the mailbox round-trip and mailbox_mu_ have
been released, preserving exactly-once callback behavior.
In `@tests/st/a2a3/host_build_graph/worker_async_fifo/test_worker_async_fifo.py`:
- Line 114: After unpacking tensors into first_a, first_b, first_out, second_a,
second_b, and second_out, explicitly delete those references before
free_host_buffer runs; apply the same cleanup at the corresponding later
unpacking block, while guarding it for early-failure paths where the names may
not be bound.
---
Outside diff comments:
In `@src/common/hierarchical/orchestrator.cpp`:
- Around line 177-206: Update fail_run_submission and cancel_unstarted_run to
synchronize phase inspection with the run/completion state, preventing
cancellation from racing scheduler activation. Recheck the current phase while
holding the appropriate lock before cancelling, and make each slot-state update
use compare-and-exchange transitions that only convert valid non-running states;
never mark an EXECUTING/RUNNING slot FAILED or overwrite a concurrent scheduler
transition.
---
Nitpick comments:
In `@python/simpler/task_interface.py`:
- Around line 1277-1284: Add concise one-line docstrings to the pipeline_depth
and runtime_slot_count properties in the task interface, matching the
neighboring properties’ style and documenting each value’s unit or semantics.
In `@src/common/hierarchical/scheduler.cpp`:
- Around line 366-369: Consolidate the repeated active-run resolution and
scoped/unscoped ready-queue branching across all four dispatch paths and the
wait predicate in the scheduler class. Add small private helpers such as
resolve_active_run() and pop_ready(...) that preserve invalid-run early returns
while centralizing queue API selection, then update the affected paths to use
them.
In `@src/common/platform/onboard/host/device_runner_base.cpp`:
- Around line 143-150: Update select_arena_bank and the associated arena
consumers to validate against the actual number of supported banks rather than
PTO_PIPELINE_MAX_DEPTH. Define or reuse a two-bank count, reject IDs outside
that range, and replace binary bank-selection logic with indexed access so each
accepted bank maps to its own arena without aliasing.
In `@src/common/platform/onboard/host/device_runner_base.h`:
- Around line 297-320: Update the documentation for record_device_orch_callable
to add an `@param` entry describing aicore_image_hash, matching its position and
purpose in the function signature. Also add the corresponding parameter
documentation to record_host_orch_callable, without changing implementation
behavior.
In `@tests/ut/cpp/hierarchical/test_orchestrator.cpp`:
- Around line 656-665: Ensure the asynchronous third begin_run test always
releases an admission slot before any ASSERT_* can exit the test, by adding a
scope guard around the first run’s slot ownership and consumption cleanup.
Preserve the existing successful-path behavior while making cleanup execute
during assertion-induced early returns, so third.get() cannot leave the test
hanging.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 976e895c-87d2-41d9-ab42-f455679ef0bc
📒 Files selected for processing (45)
docs/orchestrator.mddocs/task-flow.mddocs/worker-manager.mdpython/bindings/task_interface.cpppython/bindings/worker_bind.hpython/simpler/orchestrator.pypython/simpler/task_interface.pypython/simpler/worker.pysrc/a2a3/platform/onboard/host/device_runner.cppsrc/a2a3/platform/onboard/host/device_runner.hsrc/a2a3/runtime/host_build_graph/host/runtime_maker.cppsrc/a2a3/runtime/tensormap_and_ringbuffer/host/runtime_maker.cppsrc/a5/runtime/host_build_graph/host/runtime_maker.cppsrc/a5/runtime/tensormap_and_ringbuffer/host/runtime_maker.cppsrc/common/hierarchical/orchestrator.cppsrc/common/hierarchical/orchestrator.hsrc/common/hierarchical/scheduler.cppsrc/common/hierarchical/scheduler.hsrc/common/hierarchical/types.cppsrc/common/hierarchical/types.hsrc/common/hierarchical/worker.cppsrc/common/hierarchical/worker.hsrc/common/hierarchical/worker_manager.cppsrc/common/hierarchical/worker_manager.hsrc/common/platform/onboard/host/c_api_shared.cppsrc/common/platform/onboard/host/device_runner_base.cppsrc/common/platform/onboard/host/device_runner_base.hsrc/common/platform/sim/host/c_api_shared.cppsrc/common/task_interface/chip_callable_layout.hsrc/common/task_interface/prepare_callable_common.hsrc/common/utils/fnv1a_64.hsrc/common/worker/chip_worker.cppsrc/common/worker/chip_worker.hsrc/common/worker/pipeline_contract.hsrc/common/worker/pipeline_slot_pool.hsrc/common/worker/pto_runtime_c_api.htests/st/a2a3/host_build_graph/run_stream_reuse/kernels/aiv/kernel_sub.cpptests/st/a2a3/host_build_graph/run_stream_reuse/test_run_stream_reuse.pytests/st/a2a3/host_build_graph/worker_async_fifo/test_worker_async_fifo.pytests/ut/cpp/hierarchical/test_orchestrator.cpptests/ut/cpp/hierarchical/test_pipeline_contract.cpptests/ut/cpp/hierarchical/test_scheduler.cpptests/ut/cpp/types/test_chip_callable_upload_immutable.cpptests/ut/py/test_callable_identity.pytests/ut/py/test_worker/test_host_worker.py
ae22766 to
8fcabcf
Compare
8fcabcf to
17d945b
Compare
|
B2 was rebuilt and force-with-lease updated as a single commit on the current main. Key correction: the FIFO head can execute while its graph callback is still open, so existing L3-L2 interactive callbacks no longer deadlock. A successor still cannot dispatch until the previous run is terminal. The busy-SUB requeue path now preserves the active run partition, and unscoped queue access preserves run insertion order. Validation is recorded in the PR body. Final hardware task |
Summary
{slot_id, generation}through TaskSlot and the local chip mailbox into the B1 runtime slot.Rebase and compatibility correction
main@8e89f014; old Feat: add launch-acceptance flight fence #1467/B1 stack commits are no longer present.INVALID_RUN_ID.Validation
866 passed, 10 deselected.17d945bd: focused Python worker/L3-L2 UT184 passed.64/64 passed.task_20260728_235654_212932915886, exit 0.task_20260729_000230_235780427710, exit 0.l3_l2_message_queue, andl3_l2_orch_comm_stream.Non-goals
This PR establishes bounded whole-run admission and FIFO dispatch. The two-frame endpoint, HBG prepared epochs, and TMR prepared state remain B3, B4, and B5.